#!/usr/bin/env python3
"""
YOLO格式数据集准备脚本
将分类数据集转换为YOLO格式用于目标检测训练
"""

import os
import shutil
import yaml
from pathlib import Path
import random
from PIL import Image
import argparse

def create_yolo_annotation(image_path, class_id, output_dir):
    """为图像创建YOLO格式的标注文件"""
    # 读取图像获取尺寸
    with Image.open(image_path) as img:
        width, height = img.size
    
    # 创建YOLO格式标注（整个图像作为一个目标）
    # 格式: class_id center_x center_y width height
    center_x = 0.5  # 图像中心
    center_y = 0.5
    bbox_width = 0.9  # 边界框宽度（相对值）
    bbox_height = 0.9  # 边界框高度（相对值）
    
    # 创建标注文件
    annotation_file = output_dir / f"{image_path.stem}.txt"
    with open(annotation_file, 'w') as f:
        f.write(f"{class_id} {center_x} {center_y} {bbox_width} {bbox_height}\n")

def prepare_yolo_dataset(source_dir, output_dir, train_ratio=0.8, val_ratio=0.1):
    """准备YOLO格式数据集"""
    source_path = Path(source_dir)
    output_path = Path(output_dir)
    
    # 创建输出目录
    train_dir = output_path / "train"
    val_dir = output_path / "val"
    test_dir = output_path / "test"
    
    for dir_path in [train_dir, val_dir, test_dir]:
        dir_path.mkdir(parents=True, exist_ok=True)
        (dir_path / "images").mkdir(exist_ok=True)
        (dir_path / "labels").mkdir(exist_ok=True)
    
    # 类别映射
    class_mapping = {}
    class_id = 0
    
    # 遍历源目录
    for class_dir in source_path.iterdir():
        if class_dir.is_dir():
            class_name = class_dir.name
            class_mapping[class_name] = class_id
            
            print(f"处理类别: {class_name} (ID: {class_id})")
            
            # 获取所有图像文件
            image_files = list(class_dir.glob("*.jpg")) + list(class_dir.glob("*.jpeg")) + list(class_dir.glob("*.png"))
            random.shuffle(image_files)
            
            # 分割数据集
            total_files = len(image_files)
            train_count = int(total_files * train_ratio)
            val_count = int(total_files * val_ratio)
            
            train_files = image_files[:train_count]
            val_files = image_files[train_count:train_count + val_count]
            test_files = image_files[train_count + val_count:]
            
            # 复制文件到对应目录
            for files, target_dir in [(train_files, train_dir), (val_files, val_dir), (test_files, test_dir)]:
                for img_file in files:
                    # 复制图像
                    shutil.copy2(img_file, target_dir / "images" / img_file.name)
                    # 创建标注文件
                    create_yolo_annotation(img_file, class_id, target_dir / "labels")
            
            class_id += 1
    
    # 保存类别映射
    with open(output_path / "classes.txt", 'w') as f:
        for class_name, class_id in class_mapping.items():
            f.write(f"{class_id} {class_name}\n")
    
    # 创建数据集配置文件
    dataset_config = {
        'path': str(output_path.absolute()),
        'train': 'train/images',
        'val': 'val/images',
        'test': 'test/images',
        'nc': len(class_mapping),
        'names': list(class_mapping.keys())
    }
    
    with open(output_path / "dataset.yaml", 'w') as f:
        yaml.dump(dataset_config, f, default_flow_style=False)
    
    print(f"\n数据集准备完成！")
    print(f"输出目录: {output_path}")
    print(f"类别数量: {len(class_mapping)}")
    print(f"类别映射: {class_mapping}")
    
    return dataset_config

def main():
    parser = argparse.ArgumentParser(description="准备YOLO格式数据集")
    parser.add_argument("--source", type=str, required=True, help="源数据集目录")
    parser.add_argument("--output", type=str, required=True, help="输出目录")
    parser.add_argument("--train-ratio", type=float, default=0.8, help="训练集比例")
    parser.add_argument("--val-ratio", type=float, default=0.1, help="验证集比例")
    
    args = parser.parse_args()
    
    # 设置随机种子
    random.seed(42)
    
    # 准备数据集
    prepare_yolo_dataset(args.source, args.output, args.train_ratio, args.val_ratio)

if __name__ == "__main__":
    main()
